import torch
import numpy as np
# TODO 支持本地和url加载权重
# from torch.hub import load_state_dict_from_url
# load_state_dict_from_url(url, model_dir=None, map_location=None, progress=True, check_hash=False, file_name=None)

def load_model_weights(args):
    #------------------------------------------------------#
    #   根据预训练权重的Key和模型的Key进行加载
    #------------------------------------------------------#
    model_dict      = args.detector.state_dict()

    if args.resume:
        FLAG = True
        pretrained_dict = torch.load(args.ckpt_path, map_location = args.device)
        pretrained_dict = pretrained_dict['weights']
        print("\n\033[1;32;40m Resume weights from {} \033[0m".format(args.ckpt_path))
    elif args.model_path != None:
        FLAG = True
        pretrained_dict = torch.load(args.model_path, map_location = args.device)
        print("\n\033[1;32;40m Load weights from {} \033[0m".format(args.model_path))
    else:
        FLAG = False
    if FLAG:
        load_key, no_load_key, temp_dict = [], [], {}
        for k, v in pretrained_dict.items():
            if k in model_dict.keys() and np.shape(model_dict[k]) == np.shape(v):
                temp_dict[k] = v
                load_key.append(k)
            else:
                no_load_key.append(k)
        model_dict.update(temp_dict)

        # 检测器的预训练权重加载
        args.detector.load_state_dict(model_dict)
        #------------------------------------------------------#
        #   显示没有匹配上的Key
        #------------------------------------------------------#
        print("\nSuccessful Load Key:", str(load_key)[:500], "……\nSuccessful Load Key Num:", len(load_key))
        print("\nFail To Load Key:", str(no_load_key)[:500], "……\nFail To Load Key num:", len(no_load_key))

def weights_init(net, init_type='normal', init_gain=0.02):
    def init_func(m):
        classname = m.__class__.__name__
        if hasattr(m, 'weight') and classname.find('Conv') != -1:
            if init_type == 'normal':
                torch.nn.init.normal_(m.weight.data, 0.0, init_gain)
            elif init_type == 'xavier':
                torch.nn.init.xavier_normal_(m.weight.data, gain=init_gain)
            elif init_type == 'kaiming':
                torch.nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
            elif init_type == 'orthogonal':
                torch.nn.init.orthogonal_(m.weight.data, gain=init_gain)
            else:
                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
        elif classname.find('BatchNorm2d') != -1:
            torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
            torch.nn.init.constant_(m.bias.data, 0.0)
    print('initialize network with %s type' % init_type)
    net.apply(init_func)

def normal_init(m, mean, stddev, truncated=False):
    if truncated:
        m.weight.data.normal_().fmod_(2).mul_(stddev).add_(mean)  # not a perfect approximation
    else:
        m.weight.data.normal_(mean, stddev)
        m.bias.data.zero_()